
import {Body, Controller, Get, HttpStatus, Post, Query, Res, Headers} from '@nestjs/common';
import {response, Response} from 'express';
import {CommonRedisService} from '../service/common.redis.service';
import {CommonRequestService} from '../service/common.request.service';
import {RecordListService} from '../service/record.list.service';
import {UserListService} from '../service/user.list.service';
import {CommonResponseDto} from '../dto/common.response.dto';
import {ResCode, UserStatus} from '../common/constants';
import AppConfig from 'config/app.config';
import {formatTimeToDateStr, uuid} from "../common/utils";
import {decrypt, getModelPrice, setErrorStreamRes} from "../common/common";
import Decimal from "decimal.js";

const platform = 'volces'

@Controller('stream/' + platform)
export class VolcesController {
    constructor(private readonly commonRedisService: CommonRedisService, private readonly commonRequestService: CommonRequestService,
                private readonly recordListService: RecordListService, private readonly userListService: UserListService) {}

    @Post('/chat')
    async chat(@Headers() headers: Record<string, string>, @Body() params: {content: string, model: string }, @Res({ passthrough: false }) res: Response) {
        console.log(params)
        const modelConfig = AppConfig.ModelConfig[platform]
        if (!params.model) {
            params.model = modelConfig.dsr1.code
        }
        if (![modelConfig.dsr1.code, modelConfig.dsv3.code].includes(params.model)) {
            params.model = modelConfig.dsr1.code
        }
        const model = modelConfig[params.model].model
        const modelPrice = getModelPrice(platform, params.model)
        res.setHeader('Content-Type', 'text/event-stream; charset=utf-8');
        res.setHeader('Cache-Control', 'no-cache');
        res.setHeader('Connection', 'keep-alive');
        res.status(HttpStatus.OK);
        const headerToken = headers['api-token'] || headers['Api-Token'] || headers['ApiToken'] || headers['apiToken'] || headers['apitoken'] || ''
        if (!params.content || !headerToken) {
            setErrorStreamRes('参数缺失', res)
        } else {
            try {
                const decryptStr = decrypt(headerToken, AppConfig.ApiTokenSecretKey)
                const strArr = decryptStr.split('_')
                if (strArr.length !== 2) {
                    setErrorStreamRes('api-token非法', res)
                } else {
                    const userId: string = strArr[0]
                    const token = await this.commonRedisService.get(AppConfig.ApiHeaderKey + '_' + userId)
                    if (!token || token !== headerToken) {
                        setErrorStreamRes('非法用户', res)
                    } else {
                        const user = await this.userListService.getOneById(userId)
                        if (!user) {
                            setErrorStreamRes('用户不存在', res)
                        } else if(user.status !== UserStatus.NORMAL){
                            setErrorStreamRes('用户状态异常', res)
                        } else if(new Decimal(user.balance).lessThanOrEqualTo(0.1)){
                            setErrorStreamRes('余额不足0.1', res)
                        } else {
                            const useage = {
                                outTokens: 0,
                                inTokens: 0,
                                totalTokens: 0,
                            }
                            const response = await this.commonRequestService.postStream('https://ark.cn-beijing.volces.com/api/v3/chat/completions', {
                                model,
                                messages: [
                                    {role: 'user', content: params.content}
                                ],
                                stream: true,
                                stream_options: {include_usage:true}
                            }, {
                                headers: {
                                    Authorization: `Bearer ${AppConfig.ApiKeyConfig.volces}`
                                },
                                responseType: 'stream',
                            })
                            // @ts-ignore
                            if (response.data && typeof response.data.on === 'function') {
                                // @ts-ignore
                                response.data.on('data', (chunk) => {
                                    const messageArr = chunk.toString().replace('data: ', '').split('data: ')
                                    messageArr.forEach((msg) => {
                                        if (msg .indexOf('choices') > -1) {
                                            try {
                                                const data = JSON.parse(msg)
                                                const choices = data.choices
                                                if (data.usage) {
                                                    useage.inTokens = data.usage.prompt_tokens
                                                    useage.outTokens = data.usage.completion_tokens
                                                    useage.totalTokens = data.usage.total_tokens
                                                }
                                                if (choices.length > 1) {
                                                    choices.sort((a, b) => {
                                                        return a.index-b.index
                                                    })
                                                }
                                                let obj: any = {
                                                    id: data.Id
                                                }
                                                choices.forEach((choice) => {
                                                    const delta = choice.delta
                                                    if (delta.hasOwnProperty('reasoning_content')) {
                                                        if (!obj.reasoning) {
                                                            obj.reasoning = ''
                                                        }
                                                        obj.reasoning += delta.reasoning_content
                                                    } else if (delta.hasOwnProperty('content')) {
                                                        if (!obj.answer) {
                                                            obj.answer = ''
                                                        }
                                                        obj.answer += delta.content
                                                    }
                                                })
                                                res.write('data: '+JSON.stringify(obj)+'\n\n')
                                            }catch(e){
                                                console.log(e)
                                            }
                                        }
                                    })

                                })
                                // @ts-ignore
                                response.data.on('end', () => {
                                    console.log('本次回答计量：' + JSON.stringify(useage))
                                    const date = formatTimeToDateStr(new Date().getTime())
                                    this.recordListService.createUpdateOne({
                                        id: uuid(32),
                                        userId: userId,
                                        date,
                                        inToken: Number(useage.inTokens),
                                        outToken: Number(useage.outTokens),
                                        totalToken: Number(useage.totalTokens),
                                        model: `${platform}[#]${params.model}`,
                                        inPrice: '',
                                        outPrice: ''
                                    }, modelPrice)
                                    res.end('[DONE]')
                                })
                            }
                        }
                    }
                }
            } catch(e){
                console.log(e)
                setErrorStreamRes('服务异常', res)
            }
        }
    }
}
